#!/usr/bin/env python3
"""
Export ArangoDB data and import to Postgres with AGE support.
This is a best-effort migration script to export ObserveEntities and edges from ArangoDB
and import into Postgres with the AGE graph 'topo'.

Usage: python tools/arango2pg.py --arango-url http://127.0.0.1:8529 --arango-db gala --pg-dsn 'host=... dbname=... user=... password=...'
"""
import argparse
import json
import sys
from typing import List

from pyArango.connection import Connection

from spider.dao.arangodb import ArangoObserveEntityDaoImpl, ArangoRelationDaoImpl
from spider.dao.postgres.dao import PgObserveEntityDaoImpl, PgRelationDaoImpl
from types import SimpleNamespace
from spider.util.entity import escape_entity_id


def to_observe_entity_from_arango_doc(d: dict) -> SimpleNamespace:
    # arango docs store _key and other fields
    eid = d.get('entity_id') or d.get('_key') or d.get('id')
    if not eid and '_key' in d:
        eid = d.get('_key')
    return SimpleNamespace(id=eid,
                         type=d.get('type'),
                         level=d.get('level'),
                         attrs={k: v for k, v in d.items() if k not in ['_key', 'type', 'level', 'timestamp', 'entity_id']},
                         )


def to_relation_from_arango_doc(d: dict) -> SimpleNamespace:
    # _from is collection/key like 'ObserveEntities_1610000000/abc'
    _from = d.get('_from')
    _to = d.get('_to')
    # extract key
    sub_id = _from.split('/')[-1]
    obj_id = _to.split('/')[-1]
    sub = SimpleNamespace(id=sub_id, type=None, level=None, attrs={})
    obj = SimpleNamespace(id=obj_id, type=None, level=None, attrs={})
    return SimpleNamespace(sub_entity=sub, obj_entity=obj, type=d.get('type'), layer=d.get('layer'))


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--arango-url', required=True)
    parser.add_argument('--arango-db', required=True)
    parser.add_argument('--pg-dsn', required=True)
    parser.add_argument('--batch-size', type=int, default=1000)
    parser.add_argument('--start-ts', type=int, default=0)
    parser.add_argument('--end-ts', type=int, default=0)
    parser.add_argument('--dry-run', action='store_true')
    return parser.parse_args()


def main():
    args = parse_args()
    conn = Connection(arangoURL=args.arango_url)
    if not conn.hasDatabase(args.arango_db):
        print(f'ArangoDB {args.arango_db} not found')
        sys.exit(1)
    db = conn.databases[args.arango_db]

    # initialize PG DAO
    pg_conf = {'url': args.pg_dsn, 'minconn': 1, 'maxconn': 5, 'use_age': True}
    pg_entity_dao = PgObserveEntityDaoImpl(pg_conf)
    pg_rel_dao = PgRelationDaoImpl(pg_conf)

    # find timestamps collection
    if 'Timestamps' not in db.collections:
        print('No Timestamps collection found in Arango, nothing to migrate')
        return
    timestamps_coll = db.collections['Timestamps']
    keys = [str(doc['_key']) for doc in timestamps_coll.fetchAll()]
    keys = sorted([int(k) for k in keys])
    if args.start_ts > 0:
        keys = [k for k in keys if k >= args.start_ts]
    if args.end_ts > 0:
        keys = [k for k in keys if k <= args.end_ts]

    for ts in keys:
        print(f'Migrating timestamp {ts}')
        # observe entities
        coll_name = f'ObserveEntities_{ts}'
        if coll_name not in db.collections:
            print(f'  No collection {coll_name}, skip')
            continue
        coll = db.collections[coll_name]
        docs = coll.fetchAll()
        entities = []
        for d in docs:
            ent = to_observe_entity_from_arango_doc(d)
            ent.timestamp = ts
            entities.append(ent)
            if len(entities) >= args.batch_size:
                if not args.dry_run:
                    pg_entity_dao.add_all(ts, entities)
                entities = []
        if len(entities) > 0 and not args.dry_run:
            pg_entity_dao.add_all(ts, entities)

        # relations: iterate all collections that are edge collections (class is Edges)
        edges = []
        for name, coll_obj in db.collections.items():
            # skip timestamp and observe entities
            if name == 'Timestamps' or name.startswith('ObserveEntities'):
                continue
            # assume this is an edge collection, filter by timestamp
            try:
                docs = coll_obj.fetchAll(arangoFilter=f'timestamp=={ts}')
            except Exception:
                # not an edge collection
                continue
            for d in docs:
                r = to_relation_from_arango_doc(d)
                edges.append(r)
                if len(edges) >= args.batch_size:
                    if not args.dry_run:
                        pg_rel_dao.add_all(ts, edges)
                    edges = []
        if len(edges) > 0 and not args.dry_run:
            pg_rel_dao.add_all(ts, edges)

    print('Migration completed')


if __name__ == '__main__':
    main()

